import cv2
import numpy as np
from ultralytics import YOLO
import logging
from typing import List, Dict, Tuple, Optional

class HelmetDetector:
    """
    头盔检测器类，负责检测摩托车骑行者是否佩戴头盔
    实现思路：
    1. 从通用检测结果中提取摩托车和骑行者
    2. 建立摩托车和骑行者之间的关联
    3. 对骑行者区域执行头盔检测
    4. 输出头盔佩戴状态和违规信息
    """
    
    def __init__(self, model_path=None):
        """
        初始化头盔检测器
        
        Args:
            model_path: 头盔检测模型路径，如果为None则使用默认设置
        """
        self.logger = logging.getLogger(__name__)
        self.helmet_model = self._load_helmet_model(model_path)
        
        # 类别映射
        self.helmet_classes = {
            0: 'helmet',      # 佩戴头盔
            1: 'no_helmet'    # 未佩戴头盔
        }
        
        # 摩托车和人的类别ID（基于COCO数据集）
        self.motorbike_class_id = 3  # motorcycle
        self.person_class_id = 0     # person
        
        # 配置参数
        self.min_confidence = 0.3    # 头盔检测的最小置信度
        self.motorbike_rider_distance_threshold = 50  # 摩托车和骑行者的最大距离阈值（像素）
        
        self.logger.info("头盔检测器初始化完成")
    
    def _load_helmet_model(self, model_path):
        """
        加载头盔检测模型
        
        Args:
            model_path: 模型路径
            
        Returns:
            加载的YOLO模型
        """
        try:
            if model_path and isinstance(model_path, str):
                # 加载自定义头盔检测模型
                model = YOLO(model_path)
                self.logger.info(f"加载自定义头盔检测模型: {model_path}")
            else:
                # 如果没有提供模型路径，使用通用模型并进行后处理
                # 注意：这是临时解决方案，生产环境应该使用专门训练的头盔检测模型
                model = YOLO('yolov8n.pt')
                self.logger.warning("使用通用YOLO模型作为头盔检测模型，请尽快训练专用模型")
            return model
        except Exception as e:
            self.logger.error(f"加载头盔检测模型失败: {str(e)}")
            # 返回None，后续会使用备用检测逻辑
            return None
    
    def _extract_motorbikes_and_riders(self, detections: Dict) -> List[Dict]:
        """
        从检测结果中提取摩托车和骑行者，并建立关联
        
        Args:
            detections: 检测结果字典，包含检测到的所有对象
            
        Returns:
            包含摩托车和关联骑行者的列表
        """
        result = []
        
        # 提取摩托车和人
        motorbikes = []
        persons = []
        
        if 'boxes' in detections and detections['boxes'] is not None:
            for i, (x1, y1, x2, y2, conf, cls_id) in enumerate(zip(
                detections['boxes'].xyxy[:, 0],
                detections['boxes'].xyxy[:, 1],
                detections['boxes'].xyxy[:, 2],
                detections['boxes'].xyxy[:, 3],
                detections['boxes'].conf,
                detections['boxes'].cls
            )):
                bbox = [float(x1), float(y1), float(x2), float(y2)]
                confidence = float(conf)
                class_id = int(cls_id)
                
                # 计算中心点
                center_x = (bbox[0] + bbox[2]) / 2
                center_y = (bbox[1] + bbox[3]) / 2
                
                obj_info = {
                    'id': i,
                    'bbox': bbox,
                    'confidence': confidence,
                    'class_id': class_id,
                    'center': (center_x, center_y)
                }
                
                if class_id == self.motorbike_class_id:
                    motorbikes.append(obj_info)
                elif class_id == self.person_class_id:
                    persons.append(obj_info)
        
        # 关联摩托车和骑行者
        for motorbike in motorbikes:
            # 查找距离摩托车最近的人
            closest_rider = None
            min_distance = float('inf')
            
            for person in persons:
                # 计算距离
                distance = np.sqrt(
                    (motorbike['center'][0] - person['center'][0]) **2 +
                    (motorbike['center'][1] - person['center'][1]) **2
                )
                
                # 检查是否在合理距离内且骑行者位置合理（通常在摩托车上方或后方）
                if distance < min_distance and distance < self.motorbike_rider_distance_threshold:
                    # 进一步验证：骑行者应该在摩托车的上方或后方
                    vertical_diff = person['center'][1] - motorbike['center'][1]
                    if vertical_diff < motorbike['bbox'][3] - motorbike['bbox'][1]:
                        min_distance = distance
                        closest_rider = person
            
            # 如果找到关联的骑行者，添加到结果中
            if closest_rider:
                result.append({
                    'motorbike': motorbike,
                    'rider': closest_rider,
                    'association_score': 1.0 - (min_distance / self.motorbike_rider_distance_threshold)
                })
        
        self.logger.debug(f"提取到 {len(result)} 组摩托车和骑行者关联")
        return result
    
    def _crop_rider_region(self, frame: np.ndarray, target_info: Dict) -> np.ndarray:
        """
        裁剪骑行者区域，重点关注头部区域
        
        Args:
            frame: 原始图像帧
            target_info: 包含摩托车和骑行者信息的字典
            
        Returns:
            裁剪后的骑行者区域图像
        """
        rider = target_info['rider']
        bbox = rider['bbox']
        
        # 计算裁剪区域，重点关注上半身和头部
        x1, y1, x2, y2 = bbox
        height = y2 - y1
        
        # 调整裁剪区域，重点关注头部（上半身区域）
        # 头部通常占人体高度的1/6到1/5
        head_region_height = int(height * 0.3)  # 增加裁剪高度以确保包含头部
        head_region_y1 = max(0, int(y1))
        head_region_y2 = min(frame.shape[0], int(y1 + head_region_height * 2))  # 包含颈部和肩膀
        
        # 水平方向适当扩展，确保包含整个头部
        width = x2 - x1
        expand_factor = 0.2
        expand_width = int(width * expand_factor)
        
        head_region_x1 = max(0, int(x1 - expand_width))
        head_region_x2 = min(frame.shape[1], int(x2 + expand_width))
        
        # 裁剪图像
        cropped = frame[head_region_y1:head_region_y2, head_region_x1:head_region_x2]
        
        # 如果裁剪区域太小，返回原始骑行者区域
        if cropped.shape[0] < 20 or cropped.shape[1] < 20:
            return frame[int(y1):int(y2), int(x1):int(x2)]
        
        return cropped
    
    def _process_helmet_result(self, helmet_detection, target_info: Dict) -> Dict:
        """
        处理头盔检测结果
        
        Args:
            helmet_detection: 头盔检测原始结果
            target_info: 目标信息字典
            
        Returns:
            处理后的头盔检测结果
        """
        rider = target_info['rider']
        result = {
            'rider_id': rider['id'],
            'rider_bbox': rider['bbox'],
            'helmet_status': 'unknown',  # 默认状态
            'confidence': 0.0,
            'helmet_bbox': None,
            'violation': False,
            'detection_method': 'default'
        }
        
        # 如果使用专门的头盔检测模型
        if self.helmet_model and hasattr(helmet_detection, 'boxes') and helmet_detection.boxes is not None:
            # 查找置信度最高的检测结果
            max_conf = 0
            best_detection = None
            
            for i, (conf, cls_id) in enumerate(zip(helmet_detection.boxes.conf, helmet_detection.boxes.cls)):
                confidence = float(conf)
                class_id = int(cls_id)
                
                if confidence > max_conf and confidence >= self.min_confidence:
                    max_conf = confidence
                    best_detection = {
                        'class_id': class_id,
                        'confidence': confidence,
                        'bbox': helmet_detection.boxes.xyxy[i].tolist()
                    }
            
            # 如果有有效的检测结果
            if best_detection:
                result['helmet_status'] = self.helmet_classes.get(
                    best_detection['class_id'], 'unknown'
                )
                result['confidence'] = best_detection['confidence']
                result['helmet_bbox'] = best_detection['bbox']
                result['violation'] = (result['helmet_status'] == 'no_helmet')
                result['detection_method'] = 'model'
        else:
            # 使用基于规则的检测（作为备用方案）
            # 这里实现一个简单的基于颜色和形状的头盔检测规则
            # 注意：这只是一个临时解决方案，生产环境应使用专门训练的模型
            result = self._rule_based_helmet_detection(target_info, result)
        
        return result
    
    def _rule_based_helmet_detection(self, target_info: Dict, result: Dict) -> Dict:
        """
        基于规则的头盔检测（备用方法）
        
        Args:
            target_info: 目标信息
            result: 结果字典
            
        Returns:
            更新后的结果字典
        """
        rider_bbox = target_info['rider']['bbox']
        
        # 计算头部区域的估计位置
        head_height_ratio = 0.25  # 头部占身体比例
        head_y1 = rider_bbox[1]
        head_y2 = rider_bbox[1] + (rider_bbox[3] - rider_bbox[1]) * head_height_ratio
        
        # 基于简单规则判断
        # 注意：这只是一个示例，实际应使用更复杂的规则或深度学习模型
        # 1. 简化逻辑：如果是白天场景，暂不做判断，标记为unknown
        # 2. 在实际应用中，这里应该实现更复杂的图像处理算法
        
        result['helmet_status'] = 'unknown'
        result['confidence'] = 0.5
        result['detection_method'] = 'rule_based'
        
        # 记录使用备用方法
        self.logger.warning("使用基于规则的头盔检测，请尽快训练专用模型以提高准确率")
        
        return result
    
    def detect_helmets(self, frame: np.ndarray, detections: Dict) -> List[Dict]:
        """
        检测图像中的头盔佩戴情况
        
        Args:
            frame: 输入图像帧
            detections: 通用目标检测结果
            
        Returns:
            头盔检测结果列表
        """
        if not isinstance(frame, np.ndarray) or frame.size == 0:
            self.logger.error("无效的输入图像")
            return []
        
        if not detections:
            self.logger.warning("没有检测到任何对象")
            return []
        
        try:
            # 1. 提取摩托车和骑行者
            motorbikes_and_riders = self._extract_motorbikes_and_riders(detections)
            
            # 2. 对每个骑行者执行头盔检测
            helmet_results = []
            for target in motorbikes_and_riders:
                # 裁剪骑行者区域
                rider_roi = self._crop_rider_region(frame, target)
                
                # 执行头盔检测
                helmet_detection = None
                if self.helmet_model:
                    helmet_detection = self.helmet_model(rider_roi, conf=self.min_confidence)
                    if helmet_detection and len(helmet_detection) > 0:
                        helmet_detection = helmet_detection[0]  # 取第一个结果
                
                # 处理检测结果
                result = self._process_helmet_result(helmet_detection, target)
                
                # 添加摩托车信息
                result['motorbike_bbox'] = target['motorbike']['bbox']
                result['association_score'] = target['association_score']
                
                helmet_results.append(result)
            
            self.logger.info(f"头盔检测完成，共检测到 {len(helmet_results)} 个骑行者")
            return helmet_results
            
        except Exception as e:
            self.logger.error(f"头盔检测过程中发生错误: {str(e)}")
            return []
    
    def draw_helmet_detections(self, frame: np.ndarray, helmet_results: List[Dict]) -> np.ndarray:
        """
        在图像上绘制头盔检测结果
        
        Args:
            frame: 原始图像
            helmet_results: 头盔检测结果
            
        Returns:
            绘制后的图像
        """
        output_frame = frame.copy()
        
        for result in helmet_results:
            # 绘制骑行者边框
            rider_bbox = result['rider_bbox']
            x1, y1, x2, y2 = map(int, rider_bbox)
            
            # 根据头盔状态设置颜色
            if result['helmet_status'] == 'helmet':
                color = (0, 255, 0)  # 绿色 - 合规
                label = f"头盔: 是 ({result['confidence']:.2f})"
            elif result['helmet_status'] == 'no_helmet':
                color = (0, 0, 255)  # 红色 - 违规
                label = f"头盔: 否 ({result['confidence']:.2f})"
            else:
                color = (255, 255, 0)  # 黄色 - 未知
                label = f"头盔: 未知 ({result['confidence']:.2f})"
            
            # 绘制骑行者边框
            cv2.rectangle(output_frame, (x1, y1), (x2, y2), color, 2)
            
            # 绘制摩托车边框（使用相同颜色但透明度较低）
            if 'motorbike_bbox' in result:
                mb_bbox = result['motorbike_bbox']
                mb_x1, mb_y1, mb_x2, mb_y2 = map(int, mb_bbox)
                # 创建一个临时图像来绘制半透明边框
                overlay = output_frame.copy()
                cv2.rectangle(overlay, (mb_x1, mb_y1), (mb_x2, mb_y2), color, 1)
                alpha = 0.5  # 透明度
                cv2.addWeighted(overlay, alpha, output_frame, 1 - alpha, 0, output_frame)
            
            # 绘制标签
            cv2.putText(
                output_frame, 
                label, 
                (x1, y1 - 10), 
                cv2.FONT_HERSHEY_SIMPLEX, 
                0.5, 
                color, 
                2
            )
            
            # 如果检测到头盔，绘制头盔位置
            if result['helmet_bbox']:
                h_x1, h_y1, h_x2, h_y2 = map(int, result['helmet_bbox'])
                # 将头盔坐标映射回原始图像
                rider_x1 = rider_bbox[0]
                rider_y1 = rider_bbox[1]
                h_x1 += rider_x1
                h_y1 += rider_y1
                h_x2 += rider_x1
                h_y2 += rider_y1
                cv2.rectangle(output_frame, (h_x1, h_y1), (h_x2, h_y2), (255, 0, 0), 2)
        
        return output_frame
    
    def set_confidence_threshold(self, threshold: float):
        """
        设置头盔检测的置信度阈值
        
        Args:
            threshold: 置信度阈值 (0.0 - 1.0)
        """
        if 0.0 <= threshold <= 1.0:
            self.min_confidence = threshold
            self.logger.info(f"设置头盔检测置信度阈值: {threshold}")
        else:
            self.logger.warning(f"无效的置信度阈值: {threshold}，应在0.0到1.0之间")

# 测试代码
if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(level=logging.INFO)
    
    # 创建头盔检测器实例
    detector = HelmetDetector()
    
    print("头盔检测器初始化成功")
    print(f"当前置信度阈值: {detector.min_confidence}")
    print("注意：当前使用通用YOLO模型，建议训练专门的头盔检测模型以获得更好的检测效果")